/*
 * @Author: 姚潘涛
 * @Date: 2024-06-03 10:41:29
 * @LastEditors: 姚潘涛
 * @LastEditTime: 2024-10-10 19:17:21
 * @Description:
 *
 * Copyright (c) 2024 by pandaman, All Rights Reserved.
 */
#include "xict_common/perception/detection/yolo.h"
using namespace std;
using namespace cv;
using namespace cv::dnn;

namespace xict_common {
    namespace perception {
        namespace detection {
            /**
             * @brief 读取模型
             *
             * @param net 网络
             * @param netPath 模型路径
             * @param isCuda 是否使用cuda
             * @return true
             * @return false
             */
            bool Yolo::ReadModel(Net& net, string& netPath,
                                 bool isCuda = false) {
                try {
                    net = readNet(netPath);
                } catch (const std::exception&) {
                    return false;
                }
                // cuda
                if (isCuda) {
                    LOG_F(INFO, "use gpu");
                    net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
                    net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA_FP16);
                }
                // cpu
                else {
                    LOG_F(INFO, "use cpu");
                    net.setPreferableBackend(cv::dnn::DNN_BACKEND_DEFAULT);
                    net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
                }
                return true;
            }

/**
 * @brief 进行检测
 *
 * @param SrcImg 待检测图像
 * @param net 网络
 * @param output 检测结果图像
 * @return true 识别成功
 * @return false 识别失败
 */
#if (defined YOLOV5 && YOLOV5 == false)   // yolov7
            bool Yolo::Detect(Mat& SrcImg, Net& net, vector<Output>& output) {
                output.clear();
                Mat blob;
                int col         = SrcImg.cols;
                int row         = SrcImg.rows;
                int maxLen      = MAX(col, row);
                Mat netInputImg = SrcImg.clone();
                if (maxLen > 1.2 * col || maxLen > 1.2 * row) {
                    Mat resizeImg = Mat::zeros(maxLen, maxLen, CV_8UC3);
                    SrcImg.copyTo(resizeImg(Rect(0, 0, col, row)));
                    netInputImg = resizeImg;
                }
                vector<Ptr<Layer>> layer;
                vector<string> layer_names;
                layer_names = net.getLayerNames();
                blobFromImage(netInputImg, blob, 1 / 255.0,
                              cv::Size(netWidth, netHeight),
                              cv::Scalar(0, 0, 0), true, false);
                // 如果在其他设置没有问题的情况下但是结果偏差很大，可以尝试下用下面两句语句
                //  blobFromImage(netInputImg, blob, 1 / 255.0,
                //  cv::Size(netWidth, netHeight), cv::Scalar(104, 117, 123),
                //  true, false); blobFromImage(netInputImg, blob, 1 / 255.0,
                //  cv::Size(netWidth, netHeight), cv::Scalar(114, 114,114),
                //  true, false);
                net.setInput(blob);
                std::vector<cv::Mat> netOutputImg;
                net.forward(netOutputImg, net.getUnconnectedOutLayersNames());
#if CV_VERSION_MAJOR == 4 && CV_VERSION_MINOR == 6
                std::sort(netOutputImg.begin(), netOutputImg.end(),
                          [](Mat& A, Mat& B) {
                              return A.size[2] > B.size[2];
                          });   // opencv 4.6
#endif
                std::vector<int> classIds;   // 结果id数组
                std::vector<float> confidences;   // 结果每个id对应置信度数组
                std::vector<cv::Rect> boxes;   // 每个id矩形框
                float ratio_h = (float)netInputImg.rows / netHeight;
                float ratio_w = (float)netInputImg.cols / netWidth;
                int net_width =
                    className.size() + 5;   // 输出的网络宽度是类别数+5
                for (int stride = 0; stride < strideSize;
                     stride++) {   // stride
                    float* pdata = (float*)netOutputImg[stride].data;
                    int grid_x   = (int)(netWidth / netStride[stride]);
                    int grid_y   = (int)(netHeight / netStride[stride]);
                    for (int anchor = 0; anchor < 3; anchor++) {   // anchors
                        const float anchor_w = netAnchors[stride][anchor * 2];
                        const float anchor_h =
                            netAnchors[stride][anchor * 2 + 1];
                        for (int i = 0; i < grid_y; i++) {
                            for (int j = 0; j < grid_x; j++) {
                                float box_score = Sigmoid_x(pdata[4]);
                                ;   // 获取每一行的box框中含有某个物体的概率
                                if (box_score >= boxThreshold) {
                                    cv::Mat scores(1, className.size(),
                                                   CV_32FC1, pdata + 5);
                                    Point classIdPoint;
                                    double max_class_socre;
                                    minMaxLoc(scores, 0, &max_class_socre, 0,
                                              &classIdPoint);
                                    max_class_socre =
                                        Sigmoid_x(max_class_socre);
                                    if (max_class_socre >= classThreshold) {
                                        float x = (Sigmoid_x(pdata[0]) * 2.f -
                                                   0.5f + j) *
                                                  netStride[stride];   // x
                                        float y = (Sigmoid_x(pdata[1]) * 2.f -
                                                   0.5f + i) *
                                                  netStride[stride];   // y
                                        float w =
                                            powf(Sigmoid_x(pdata[2]) * 2.f,
                                                 2.f) *
                                            anchor_w;   // w
                                        float h =
                                            powf(Sigmoid_x(pdata[3]) * 2.f,
                                                 2.f) *
                                            anchor_h;   // h
                                        int left =
                                            (int)(x - 0.5 * w) * ratio_w + 0.5;
                                        int top =
                                            (int)(y - 0.5 * h) * ratio_h + 0.5;
                                        classIds.push_back(classIdPoint.x);
                                        confidences.push_back(max_class_socre *
                                                              box_score);
                                        boxes.push_back(Rect(left, top,
                                                             int(w * ratio_w),
                                                             int(h * ratio_h)));
                                    }
                                }
                                pdata += net_width;   // 下一行
                            }
                        }
                    }
                }

                // 执行非最大抑制以消除具有较低置信度的冗余重叠框（NMS）
                vector<int> nms_result;
                NMSBoxes(boxes, confidences, nmsScoreThreshold, nmsThreshold,
                         nms_result);
                for (int i = 0; i < nms_result.size(); i++) {
                    int idx = nms_result[i];
                    Output result;
                    result.id         = classIds[idx];
                    result.confidence = confidences[idx];
                    result.box        = boxes[idx];
                    output.push_back(result);
                }
                if (output.size())
                    return true;
                else
                    return false;
            }
#else
            // yolov5
            bool Yolo::Detect(Mat& SrcImg, Net& net, vector<Output>& output) {
                output.clear();
                Mat blob;
                int col         = SrcImg.cols;
                int row         = SrcImg.rows;
                int maxLen      = max(col, row);
                Mat netInputImg = SrcImg.clone();
                if (maxLen > 1.2 * col || maxLen > 1.2 * row) {
                    Mat resizeImg = Mat::zeros(maxLen, maxLen, CV_8UC3);
                    SrcImg.copyTo(resizeImg(Rect(0, 0, col, row)));
                    netInputImg = resizeImg;
                }
                vector<Ptr<Layer>> layer;
                vector<string> layer_names;
                layer_names = net.getLayerNames();
                blobFromImage(netInputImg, blob, 1 / 255.0,
                              cv::Size(netWidth, netHeight),
                              cv::Scalar(0, 0, 0), true, false);
                // 如果在其他设置没有问题的情况下但是结果偏差很大，可以尝试下用下面两句语句
                //  blobFromImage(netInputImg, blob, 1 / 255.0,
                //  cv::Size(netWidth, netHeight), cv::Scalar(104, 117, 123),
                //  true, false); blobFromImage(netInputImg, blob, 1 / 255.0,
                //  cv::Size(netWidth, netHeight), cv::Scalar(114, 114,114),
                //  true, false);
                net.setInput(blob);
                std::vector<cv::Mat> netOutputImg;
                net.forward(netOutputImg, net.getUnconnectedOutLayersNames());
                std::vector<int> classIds;   // 结果id数组
                std::vector<float> confidences;   // 结果每个id对应置信度数组
                std::vector<cv::Rect> boxes;   // 每个id矩形框
                float ratio_h = (float)netInputImg.rows / netHeight;
                float ratio_w = (float)netInputImg.cols / netWidth;
                int net_width =
                    className.size() + 5;   // 输出的网络宽度是类别数+5
                float* pdata = (float*)netOutputImg[0].data;
                for (int stride = 0; stride < strideSize;
                     stride++) {   // stride

                    int grid_x = (int)(netWidth / netStride[stride]);
                    int grid_y = (int)(netHeight / netStride[stride]);
                    for (int anchor = 0; anchor < 3; anchor++) {   // anchors
                        // const float anchor_w = netAnchors[stride][anchor *
                        // 2]; const float anchor_h = netAnchors[stride][anchor
                        // * 2 + 1];
                        for (int i = 0; i < grid_y; i++) {
                            for (int j = 0; j < grid_x; j++) {
                                float box_score = pdata[4];
                                ;   // 获取每一行的box框中含有某个物体的概率
                                if (box_score >= boxThreshold) {
                                    cv::Mat scores(1, className.size(),
                                                   CV_32FC1, pdata + 5);
                                    Point classIdPoint;
                                    double max_class_socre;
                                    minMaxLoc(scores, 0, &max_class_socre, 0,
                                              &classIdPoint);
                                    max_class_socre = max_class_socre;
                                    if (max_class_socre >= classThreshold) {
                                        float x = pdata[0];   // (Sigmoid_x(pdata[0])
                                                              // * 2.f - 0.5f
                                                              // + j) *
                                                              // m_Mark_Stride[stride];
                                                              // //x
                                        float y = pdata[1];   //(Sigmoid_x(pdata[1])
                                                              //* 2.f - 0.5f
                                                              //+
                                                              // i) *
                                                              // m_Mark_Stride[stride];
                                                              // //y
                                        float w =
                                            pdata[2];   // powf(Sigmoid_x(pdata[2])
                                                        // * 2.f, 2.f) *
                                                        // anchor_w*ratio_c; //w
                                        float h =
                                            pdata[3];   // powf(Sigmoid_x(pdata[3])
                                                        // * 2.f, 2.f) *
                                                        // anchor_h*ratio_r; //h
                                        int left =
                                            round((x - 0.5 * w) * ratio_w);
                                        int top =
                                            round((y - 0.5 * h) * ratio_h);
                                        classIds.push_back(classIdPoint.x);
                                        confidences.push_back(max_class_socre *
                                                              box_score);
                                        boxes.push_back(
                                            Rect(left, top, round(w * ratio_w),
                                                 round(h * ratio_h)));
                                    }
                                }
                                pdata += net_width;   // 下一行
                            }
                        }
                    }
                }

                // 执行非最大抑制以消除具有较低置信度的冗余重叠框（NMS）
                vector<int> nms_result;
                NMSBoxes(boxes, confidences, nmsScoreThreshold, nmsThreshold,
                         nms_result);
                for (int i = 0; i < nms_result.size(); i++) {
                    int idx = nms_result[i];
                    Output result;
                    result.id         = classIds[idx];
                    result.confidence = confidences[idx];
                    result.box        = boxes[idx];
                    output.push_back(result);
                }
                if (output.size())
                    return true;
                else
                    return false;
            }
#endif

            /**
             * @brief 判断两个长方形是否包含
             *
             * @param rect1
             * @param rect2
             * @return true
             * @return false
             */
            bool isInside(Rect rect1, Rect rect2) {
                return (rect1 == (rect1 & rect2));
            }

            /**
             * @brief 绘制出识别结果
             *
             * @param img 原始图像
             * @param result 识别结果
             * @param color 绘制颜色
             * @return cv::Mat 包含识别结果的图像
             */
            cv::Mat Yolo::DrawPred(Mat& img, vector<Output> result,
                                   vector<Scalar> color) {
                for (size_t i = 0; i < result.size(); i++) {
                    int left, top;
                    left          = result[i].box.x;
                    top           = result[i].box.y;
                    int color_num = i;
                    if (className[result[i].id] != "person" and
                        className[result[i].id] != "car")
                        continue;

                    // 显示bbox
                    rectangle(img, result[i].box, color[result[i].id], 2, 8);

                    // 显示类别
                    // string label =
                    //     className[result[i].id] + ":" +
                    //     to_string(result[i].confidence);
                    // int baseLine;
                    // Size labelSize =
                    //     getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1,
                    //     &baseLine);
                    // top = max(top, labelSize.height);
                    // putText(img, label, Point(left, top),
                    // FONT_HERSHEY_SIMPLEX, 1,
                    //         color[result[i].id], 2);

                    // 显示logo
                    // Mat imageROI;
                    // string img_path = "../image/" + className[result[i].id] +
                    // ".png"; Mat logo = imread(img_path); if (logo.empty()) {
                    //     LOG_F(INFO, "%s is empty", img_path.c_str());
                    // }

                    // Rect rect = Rect(0, 0, img.cols, img.rows);
                    // if (isInside(Rect(left + int(result[i].box.width * 0.5),
                    // top - logo.rows, logo.cols, logo.rows), rect)) {
                    //     imageROI = img(Rect(left + int(result[i].box.width *
                    //     0.5), top - logo.rows, logo.cols, logo.rows));
                    //     resize(logo, logo, imageROI.size());
                    //     addWeighted(imageROI, 0.5, logo, 1, 0, imageROI);
                    // }
                    // Rect rect = Rect(0, 0, img.cols, img.rows);
                    // if (isInside(Rect(left, top, result[i].box.width,
                    // result[i].box.height), rect)){
                    //     imageROI =
                    //         img(Rect(left, top, result[i].box.width,
                    //         result[i].box.height));
                    //     resize(logo, logo, imageROI.size());
                    //     addWeighted(imageROI, 1, logo, 1, 0, imageROI);
                    // }
                }
                return img;
            }
        }   // namespace detection
    }       // namespace perception
}   // namespace xict_common